package com.dev.toolkit.config.ds;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.transaction.annotation.EnableTransactionManagement;

import javax.sql.DataSource;
import java.util.HashMap;
import java.util.Map;

/**
 * @author ZhuYouBin
 * @version 1.0.0
 * @date 2024/5/21 11:35
 * @description 动态数据源
 */
// 开启事务管理
@EnableTransactionManagement
public class DynamicDataSource extends AbstractRoutingDataSource {
    private static final Logger logger = LoggerFactory.getLogger(DynamicDataSource.class);
    /**
     * 定义目标数据源
     */
    public static final Map<Object, Object> DEFINE_TARGET_DATA_SOURCE_MAP = new HashMap<>();

    /**
     * 构造方法
     *
     * @param defaultDataSource   默认数据源
     * @param targetDataSourceMap 所有的目标数据源
     */
    public DynamicDataSource(DataSource defaultDataSource, Map<Object, Object> targetDataSourceMap) {
        // 设置默认的数据源对象
        super.setDefaultTargetDataSource(defaultDataSource);
        // 设置所有的目标数据源对象
        super.setTargetDataSources(targetDataSourceMap);
        // 备份目标数据源
        DEFINE_TARGET_DATA_SOURCE_MAP.putAll(targetDataSourceMap);
    }

    /**
     * 决定当前线程要使用哪个数据源key
     */
    @Override
    protected Object determineCurrentLookupKey() {
        // 获取当前线程所使用的数据源名称
        String dataSourceType = DsContextHolder.getDataSourceType();
        logger.info("当前使用的数据源是:[{}]", dataSourceType);
        return dataSourceType;
    }

    /**
     * 添加新的数据源，并且需要重新调用一下 afterPropertiesSet() 初始化回调方法，用于更新目标数据源
     *
     * @param dataSourceName 数据源名称
     * @param dataSource     数据源对象
     */
    public void addDataSource(String dataSourceName, DataSource dataSource) {
        // 保存数据源
        DEFINE_TARGET_DATA_SOURCE_MAP.put(dataSourceName, dataSource);
        // 更新目标数据源
        super.setTargetDataSources(DEFINE_TARGET_DATA_SOURCE_MAP);
        // 调用回调方法，重新设置数据源
        super.afterPropertiesSet();
    }

    /**
     * 是否已经包含这个数据源名称
     *
     * @param dataSourceName 数据源名称
     * @return 包含当前数据源则返回true
     */
    public static boolean containsDataSource(String dataSourceName) {
        return DEFINE_TARGET_DATA_SOURCE_MAP.containsKey(dataSourceName);
    }
}